#include <stdlib.h>


#include "kvec.h"
#include "utils.h"
#include "bamMIOSortCore.h"


//  Step1: Read multi-bams , output multi-bams(buck)


bam_hdr_t *h_ ; 

long buckSize = 0 ;

uint64_t num_bps;

int numOutput ;

samFile **fpn  = NULL ;

long *targetOffset_ = NULL ; 

int  numTarget_ =  0 ;

#define BUFFERSIZE  1024*1024*4

typedef kvec_t(bam1_t) bamVect ;

typedef kvec_t(char *) strVect ;


char *temp_fileName =  NULL ;

#define STATUS_INIT 0 
#define STATUS_RUNNING 1
#define STATUS_END 2



typedef struct {
    int offset ;
    samFile *fout ;
    int n_threads ;
} ktp_aux_t ;



typedef struct {

    bamVect   bamVect_ ;

} ktp_data_t ;



void kt_pipeline(int n_threads, void *(*func)(void*, int, void*), void *shared_data, int n_steps);


inline long  offset2coordinate( int tid , int pos)
{
    return  tid >=0 ? targetOffset_[tid] + pos : num_bps ;
}

#if 0
static  void *process_classify(void *shared , int step , void *_data)
{
    ktp_step1_axu_t *aux = (ktp_step1_axu_t *) shared ; 
    ktp_step1_data_t *data= (ktp_step1_data_t *) _data ;
    int  i =  0  ;
    
    if(step == 0){
        ktp_step1_data_t  *ret  =   calloc(1,sizeof(ktp_step1_data_t));

        kv_init(ret->bamVect_);
        kv_resize(bam1_t , ret->bamVect_ , BUFFERSIZE);
        kv_setALL(bam1_t, ret->bamVect_);

        if( aux->status == STATUS_END && aux->offset ==  kv_size(aux->strVect_)){
            bam_hdr_destroy(aux->h_);
            sam_close(aux->fp);
            kv_destroy(ret->bamVect_);
            free(ret);
            return NULL ;
        }else {

            if(aux->status == STATUS_END || aux->status == STATUS_INIT){
                char *item = aux->strVect_.a[aux->offset] ;
                fprintf(stderr,"file Name = %s\n",item);
                //  Close  previous opened File description 
                if( aux->status != STATUS_INIT){
                    bam_hdr_destroy(aux->h_);
                    sam_close(aux->fp);
                }
                aux->fp = sam_open(item,"r");
                if(!aux->fp)  {
                    bam_hdr_destroy(aux->h_);
                    kv_destroy(ret->bamVect_);
                    free(ret);
                    return  NULL ;
                }
                aux->h_ = sam_hdr_read(aux->fp);
            }
            for( i = 0 ; i < kv_max(ret->bamVect_) ;i++){
                bam1_t  *bam =  &ret->bamVect_.a[i] ;
                if( sam_read1(aux->fp,aux->h_,bam) < 0 ) {
                    break ; 
                }

            }
            ret->bamVect_.n =  i ;
            if( i != kv_max(ret->bamVect_)){
                aux->offset++ ;
                aux->status = STATUS_END ;
            }else {
                fprintf(stderr,"%zu\n",kv_size(ret->bamVect_));
                aux->status = STATUS_RUNNING ;
            }
	   
            return ret ;
        }


    }else if (step == 1){

        fprintf(stderr,"%zu\n",kv_size(data->bamVect_));

       for( i = 0 ; i < kv_size(data->bamVect_) ; i++){
            bam1_t *bam = &data->bamVect_.a[i];
            long offset_ = offset2coordinate(bam->core.tid,bam->core.pos);
            int  id = offset_ / buckSize ;
            if(sam_write1(aux->fps[id],aux->h_,bam) < 0){

                fprintf(stderr,"Error:Write Bam\n");

            }
            free(bam->data);
        } 

        kv_destroy(data->bamVect_);
        free(data);
        return 0 ;
    }
    return  0 ;
}
#endif
void coreInit( char *fn , char *temp , int n )
{

    samFile  *fp = sam_open(fn,"r");

    int i = 0 ;

    numOutput = n ;

    h_ = sam_hdr_read(fp);

    temp_fileName = temp ;

    char *outputName = malloc(sizeof(char)*( strlen(temp) + 40));

    fpn = malloc(sizeof(samFile *)*numOutput);

    for( i = 0 ; i <  numOutput ; i++){
        sprintf(outputName ,"%s.%d.bam",temp,i);
        fprintf(stderr,"outputName=%s\n",outputName);
        fpn[i] = sam_open(outputName,"wb");
        hts_set_threads(fpn[i],48);
        if(sam_hdr_write(fpn[i],h_) < 0){
            break ;
        }
        memset(outputName , 0 , sizeof(char)*(strlen(temp)+32));
    }
    
    for ( i = 0 ; i  < h_->n_targets ; i++)
        num_bps += h_->target_len[i];

    buckSize = (long) (((double) (num_bps + 2 ) / (double)n) + 0.5);

    targetOffset_ =  calloc(h_->n_targets,sizeof(long));
    targetOffset_[0] = 0 ; 

    for( i = 1 ;  i  <  h_->n_targets ; i++){
        targetOffset_[i] +=  (h_->target_len[i-1] +  targetOffset_[i-1]);
    } 
    numTarget_ = h_->n_targets ;
    free(outputName);
    sam_close(fp);
}

void coreDestory()
{
        bam_hdr_destroy(h_);
}


void run_classify_buck( char **fns , int n  , int n_threads)
{

    int  i = 0 ;

    bam1_t *b = bam_init1();

    for ( i = 0 ; i < n ; i++){

        double start = get_run_time();

        samFile *fp = sam_open(fns[i],"r");
        hts_set_threads(fp,n_threads);
        bam_hdr_t *h =  sam_hdr_read(fp);
        
        if(!fp){
            fprintf(stderr,"Can't open %s'\n",fns[i]);
            break ;
        }

        fprintf(stderr,"FileName %s: time = \t",fns[i]);

        do{

            if( sam_read1(fp,h,b) < 0){
                break ;
            }

            long offset_ = offset2coordinate(b->core.tid,b->core.pos);

            int  id = offset_ / buckSize ;

            if(sam_write1(fpn[id],h,b) < 0){
                fprintf(stderr,"Write Error\n");
                break ;
            }
    
        }while(1);
        fprintf(stderr,"%.2f sec\n",get_run_time() - start );
        sam_close(fp);
        bam_hdr_destroy(h);
    }

    bam_destroy1(b);

    for( i = 0 ; i < numOutput ; i++){
        sam_close(fpn[i]);
    }

    free(fpn);

}


int compared(const void *a_ , const void *b_)
{
    const bam1_t *a = a_ ;
    const bam1_t *b = b_ ;
    return  (((uint64_t)a->core.tid<<32|(a->core.pos+1)<<1|bam_is_rev(a)) > ((uint64_t)b->core.tid<<32|(b->core.pos+1)<<1|bam_is_rev(b))); 

}

static void *process( void *shared , int step , void *_data)
{

    ktp_aux_t *aux =  ( ktp_aux_t *) shared;
    ktp_data_t *data = ( ktp_data_t *) _data ;

    if(step == 0){
        if(aux->offset == numOutput){
            return NULL ;
        }
        ktp_data_t *ret =  calloc(1,sizeof(ktp_data_t));
        kv_resize(bam1_t ,ret->bamVect_ , BUFFERSIZE);
        kv_setALL(bam1_t ,ret->bamVect_);

        char *inputFile = malloc(sizeof(char)*( strlen(temp_fileName) + 40));

        sprintf(inputFile,"%s.%d.bam",temp_fileName,aux->offset);

        fprintf(stderr,"Read:%s\n",inputFile);
        samFile *fp =  sam_open( inputFile, "r");
        bam_hdr_t *h = sam_hdr_read(fp);
        hts_set_threads(fp,24);

        do{

            bam1_t  *bam =  &ret->bamVect_.a[ret->bamVect_.n] ;
            if( sam_read1(fp,h,bam) < 0 ) {
                break ; 
            }
            ret->bamVect_.n++ ;
            if( ret->bamVect_.m == ret->bamVect_.n ){
                ret->bamVect_.m = ret->bamVect_.m << 1 ;
                ret->bamVect_.a =  realloc(ret->bamVect_.a ,  sizeof(bam1_t)*ret->bamVect_.m);
                memset(ret->bamVect_.a+ret->bamVect_.n,0,sizeof(bam1_t)*(ret->bamVect_.m-ret->bamVect_.n));
            }

        }while(1);

        free(inputFile);
        sam_close(fp);
        aux->offset++ ;
        return ret ;

    }else if ( step == 1){
        qsort(data->bamVect_.a, kv_size(data->bamVect_) , sizeof(bam1_t) , compared);

        return data ;
    }else if ( step == 2){
        int i =  0 ;
        for( i = 0 ; i <  kv_size(data->bamVect_) ; i++){
            bam1_t  *bam = &data->bamVect_.a[i] ;

            if(  sam_write1(aux->fout,h_,bam) < 0){
                fprintf(stderr,"[Error] Write bam Failed\n");
                break ;
            }
            free(bam->data);
        }
        kv_destroy(data->bamVect_);
    }

    return NULL ;
}

void run_sort_buck( char *fn , int n_threads)
{

    ktp_aux_t  *aux = calloc(1,sizeof(ktp_aux_t));


    aux->fout = sam_open(fn,"wb");

    hts_set_threads(aux->fout , n_threads) ;

   if( sam_hdr_write( aux->fout , h_) < 0 ){

        fprintf(stderr,"[Error] Header Failed\n");
    }

    aux->n_threads = n_threads ;

//    aux->status = STATUS_INIT ;

    kt_pipeline(3,process,aux,3);

    sam_close(aux->fout);
    free(aux);
#if 0
    ktp_step1_axu_t *aux = calloc(1,sizeof(ktp_step1_axu_t));

    aux->status = STATUS_INIT ;

    aux->fps = fpn ;

    aux->offset = 0 ;

    aux->n_threads = n_threads ; 

    int i = 0 ; 

    for( i = 0 ; i  < n ; i++){
        // Push bams FileName  into vector
        kv_push( char * , aux->strVect_ , fns[i]);
    }

    kt_pipeline(2,process_classify,aux,2);



    for( i = 0 ; i < numOutput ; i++){
        sam_close(aux->fps[i]);
    }

    
    bam_hdr_destroy(h_);

    free(aux->fps);
    kv_destroy(aux->strVect_);
    free(aux);
#endif
    

}




